# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: leon

import datetime
import logging
import json
from typing import List, Union

from tornado import gen

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.base.timestamp import get_timestamp
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase
from coin.exchange.okex.parse_util import parse_okex_msg3
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.okex.kr_rest.product import OkexProduct


class OkexV3SubscriberBase(WebsocketJsonSubscriberBase):
  BookFeedDisconnectWindowInNs = 2 * 1e9  # 2 sec

  _subscriber_channel = None
  _flow_topic_prefix = None

  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None,
               *,
               subscribe_ticker=False,
               subscribe_price_range=True,
               subscribe_funding_rate=False,
               subscribe_mark_price=False,
               need_monitor_feed=False):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._symbol = to_list(symbol)
    if self._subscriber_channel == 'swap':
      # Native symbols are given, do nothing.
      pass
    elif self._subscriber_channel == 'futures':
      # v1 native symbols are given, transfer to v3 native symbols.
      # Before call FromNativeSubscriptionSymbolV1, current_datetime
      # need to be specific to avoid cache mistake.
      # New added USDT futures use v3 native symbols IN UPPER CASE.
      current_datetime = datetime.datetime.now()
      self._symbol = [
          OkexFuturesProduct.FromNativeSubscriptionSymbolV1(
              symbol, current_datetime=current_datetime).native_symbol_v3
          if symbol.islower() else symbol for symbol in self._symbol
      ]
    elif self._subscriber_channel == 'spot':
      # v1 native symbols are given, transfer to v3 native symbols.
      self._symbol = [
          OkexProduct.FromStrNativeProductV1(symbol).native_symbol_v3 for symbol in self._symbol
      ]
    else:
      raise ValueError(self._subscriber_channel)
    self._register_url('wss://real.okex.com:8443/ws/v3')

    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)

    self._last_pong_time = None
    self._subscribe_ticker = subscribe_ticker
    self._subscribe_price_range = subscribe_price_range
    self._subscribe_funding_rate = subscribe_funding_rate
    self._subscribe_mark_price = subscribe_mark_price

    self._need_monitor_feed = need_monitor_feed
    self._last_book_ts = None
    self._monitor_handle = None

  @gen.coroutine
  def on_open(self):
    self._logger.debug('Connected')
    self._ioloop.add_callback(self._send_subscribe_request)

  def _decode(self, ws_msg):
    return parse_okex_msg3(ws_msg)

  @gen.coroutine
  def ping(self):
    yield self._ws.write_message('ping')
    self._logger.debug('ping')

  @gen.coroutine
  def _send_subscribe_request(self):
    for symbol in self._symbol:
      self._logger.debug('Subscribing %s' % symbol)

      yield self._ws.write_message(
          json.dumps({
              'op': 'subscribe', 'args': ['%s/depth:%s' % (self._subscriber_channel, symbol)]
          }))

      if self._subscriber_channel == 'futures':
        yield self._ws.write_message(
            json.dumps({
                'op': 'subscribe',
                'args': ['%s/depth_l2_tbt:%s' % (self._subscriber_channel, symbol)]
            }))

      yield self._ws.write_message(
          json.dumps({
              'op': 'subscribe', 'args': ['%s/trade:%s' % (self._subscriber_channel, symbol)]
          }))

      if self._subscribe_ticker:
        yield self._ws.write_message(
            json.dumps({
                'op': 'subscribe', 'args': ['%s/ticker:%s' % (self._subscriber_channel, symbol)]
            }))

      if self._subscribe_price_range and self._subscriber_channel != 'spot':
        yield self._ws.write_message(
            json.dumps({
                'op': 'subscribe',
                'args': ['%s/price_range:%s' % (self._subscriber_channel, symbol)]
            }))

      if self._subscribe_funding_rate and self._subscriber_channel == 'swap':
        yield self._ws.write_message(
            json.dumps({
                'op': 'subscribe',
                'args': ['%s/funding_rate:%s' % (self._subscriber_channel, symbol)]
            }))

      if self._subscribe_mark_price and self._subscriber_channel != 'spot':
        yield self._ws.write_message(
            json.dumps({
                'op': 'subscribe',
                'args': ['%s/mark_price:%s' % (self._subscriber_channel, symbol)]
            }))

      yield gen.sleep(0.5)  # Sleep 0.5 second before next request

    self._publish_topic_event('START')
    self._logger.info('Subscribe done')

    self._last_pong_time = datetime.datetime.now()
    while self._state == self.State.RUNNING:
      if (datetime.datetime.now() - self._last_pong_time >= datetime.timedelta(minutes=2)):
        self._logger.error('No pong has been received since %s. Closing...'
                           % str(self._last_pong_time))
        self._notify_close()
        break

      self.ping()
      yield gen.sleep(10)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_depth_fullbook = topic_map.search_topic_id(
        '%s_depth_fullbook_%s%s' % (self._flow_topic_prefix, group_id, worker_id))[0]
    topic_depth_diff = topic_map.search_topic_id('%s_depth_diff_%s%s' %
                                                 (self._flow_topic_prefix, group_id, worker_id))[0]
    topic_trade = topic_map.search_topic_id('%s_trade_%s%s' %
                                            (self._flow_topic_prefix, group_id, worker_id))[0]
    topic_ticker = topic_map.search_topic_id('%s_ticker_%s%s' %
                                             (self._flow_topic_prefix, group_id, worker_id))[0]

    self._publisher = {
        'depth_fullbook': queue.get_writer(topic_id=topic_depth_fullbook[1]),
        'depth_diff': queue.get_writer(topic_id=topic_depth_diff[1]),
        'trade': queue.get_writer(topic_id=topic_trade[1]),
        'ticker': queue.get_writer(topic_id=topic_ticker[1])
    }

    target_topic_list = [
        topic_depth_fullbook[0], topic_depth_diff[0], topic_trade[0], topic_ticker[0]
    ]

    if self._subscriber_channel == 'spot':
      pass
    elif self._subscriber_channel == 'futures':
      topic_price_range = topic_map.search_topic_id(
          '%s_price_range_%s%s' % (self._flow_topic_prefix, group_id, worker_id))[0]
      topic_mark_price = topic_map.search_topic_id(
          '%s_mark_price_%s%s' % (self._flow_topic_prefix, group_id, worker_id))[0]
      topic_realtime_depth_fullbook = topic_map.search_topic_id(
          '%s_realtime_depth_fullbook_%s%s' % (self._flow_topic_prefix, group_id, worker_id))[0]
      topic_realtime_depth_diff = topic_map.search_topic_id(
          '%s_realtime_depth_diff_%s%s' % (self._flow_topic_prefix, group_id, worker_id))[0]
      self._publisher['price_range'] = queue.get_writer(topic_id=topic_price_range[1])
      self._publisher['mark_price'] = queue.get_writer(topic_id=topic_mark_price[1])
      self._publisher['realtime_depth_fullbook'] = queue.get_writer(
          topic_id=topic_realtime_depth_fullbook[1])
      self._publisher['realtime_depth_diff'] = queue.get_writer(
          topic_id=topic_realtime_depth_diff[1])
      target_topic_list.extend([
          topic_price_range[0],
          topic_mark_price[0],
          topic_realtime_depth_fullbook[0],
          topic_realtime_depth_diff[0]
      ])
    elif self._subscriber_channel == 'swap':
      topic_price_range = topic_map.search_topic_id(
          '%s_price_range_%s%s' % (self._flow_topic_prefix, group_id, worker_id))[0]
      topic_funding_rate = topic_map.search_topic_id(
          '%s_funding_rate_%s%s' % (self._flow_topic_prefix, group_id, worker_id))[0]
      topic_mark_price = topic_map.search_topic_id(
          '%s_mark_price_%s%s' % (self._flow_topic_prefix, group_id, worker_id))[0]
      self._publisher['price_range'] = queue.get_writer(topic_id=topic_price_range[1])
      self._publisher['funding_rate'] = queue.get_writer(topic_id=topic_funding_rate[1])
      self._publisher['mark_price'] = queue.get_writer(topic_id=topic_mark_price[1])
      target_topic_list.extend([topic_price_range[0], topic_funding_rate[0], topic_mark_price[0]])
    else:
      raise ValueError(self._subscriber_channel)

    self._enable_topic_event_publisher(queue, 'okex_ws_topic_event', target_topic_list)

    self._add_teardown_callback(self._teardown_publisher)

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def _on_subscribe(self, timestamp: int, message: dict):
    try:
      channel = message['channel']
      self._logger.debug('Channel subscribed: %s' % channel)
    except KeyError:
      self._logger.error('Invalid message\n%s' % message)

  def _on_depth(self, timestamp: int, message: dict):
    action = message['action']
    if action == 'partial':
      self._publisher['depth_fullbook'].write(timestamp, message)
    elif action == 'update':
      self._publisher['depth_diff'].write(timestamp, message)
    else:
      raise ValueError(message)

  def _on_realtime_depth(self, timestamp: int, message: dict):
    action = message['action']
    if action == 'partial':
      self._publisher['realtime_depth_fullbook'].write(timestamp, message)
    elif action == 'update':
      self._publisher['realtime_depth_diff'].write(timestamp, message)
    else:
      raise ValueError(message)

  def on_message(self, timestamp: int, message: Union[dict, str]):
    if self._state != self.State.RUNNING:
      return

    if message is None:
      return

    if isinstance(message, str):
      assert message == 'pong'
      self._last_pong_time = datetime.datetime.now()
      self._logger.debug('pong')
    elif isinstance(message, dict):
      event = message.get('event', None)
      table = message.get('table', None)
      assert event is None or table is None

      if event == 'subscribe':
        self._on_subscribe(timestamp, message)
      elif table == '%s/depth' % self._subscriber_channel:
        self._on_depth(timestamp, message)
        self._last_book_ts = timestamp
      elif table == '%s/depth_l2_tbt' % self._subscriber_channel:
        self._on_realtime_depth(timestamp, message)
        self._last_book_ts = timestamp
      elif table == '%s/trade' % self._subscriber_channel:
        if self._last_pong_time is not None:
          self._publisher['trade'].write(timestamp, message)
      elif table == '%s/ticker' % self._subscriber_channel:
        self._publisher['ticker'].write(timestamp, message)
      elif table == '%s/price_range' % self._subscriber_channel:
        self._publisher['price_range'].write(timestamp, message)
      elif table == '%s/funding_rate' % self._subscriber_channel:
        self._publisher['funding_rate'].write(timestamp, message)
      elif table == '%s/mark_price' % self._subscriber_channel:
        self._publisher['mark_price'].write(timestamp, message)
      else:
        self._logger.error('Unknown message: %s' % str(message))
    else:
      raise ValueError(message)

  def _monitor_feed(self):
    self._monitor_handle = self._ioloop.add_timeout(
        datetime.timedelta(seconds=self.BookFeedDisconnectWindowInNs), self._monitor_feed)

    if self._last_book_ts is None:
      return

    diff = get_timestamp() - self._last_book_ts
    if diff < self.BookFeedDisconnectWindowInNs:
      return

    self._logger.error('Monitor feed detect discontinuity.')
    self._notify_close()

  def start(self, ioloop, delay_time=None):
    super().start(ioloop, delay_time)
    if self._need_monitor_feed:
      self._monitor_feed()

  def stop(self, trigger_close_callback=False):
    if self._need_monitor_feed and self._monitor_handle is not None:
      self._ioloop.remove_timeout(self._monitor_handle)
    super().stop(trigger_close_callback)
